import time

import streamlit as st

from api.demo_api import Tongyi_llm
from rag.file_rag import get_response_from_file

st.set_page_config(layout="wide")

st.title("CLOSEAI")

# 定义布局（侧边栏）
with st.sidebar:
    #  定义标题
    st.title("CLOSEAI")
    st.divider()
    # 通义⼤模型密钥输⼊
    key = st.text_input("请输⼊通义大模型api key：", type="password")
    # key获取跳转链接
    st.page_link(
        "https://bailian.console.aliyun.com/?spm=5176.28326591.0.0.40f76ee1vUR5Of&accounttraceid=d873c9fe82ea443d9f3170f30dd5690cfnng#/api_key",
        label="⽆key？请点我获取", icon="♿")

    st.divider()

    # 在侧边栏中添加链的类型选择
    chain_choice = st.selectbox(
        '选择链的类型',
        ('使用新版create_retriever_chain', '使用老版ConversationalRetrievalChain')
    )

    # 当选择使用老版ConversationalRetrievalChain时，展示策略选择
    if chain_choice == '使用老版ConversationalRetrievalChain':
        strategy_choice = st.selectbox(
            '选择ConversationalRetrievalChain的策略',
            ('stuff', 'map_reduce', 'refine', 'map_rerank')
        )
    else:
        strategy_choice = None

    st.session_state["chain_choice"] = chain_choice
    st.session_state["strategy_choice"] = strategy_choice

# 创建一个可折叠的内容区域
expander = st.expander('选择模型')
expander.write('只有通义千问')


# 定义⼀个模拟流式输出的⽅法
def stream_data(result_text):
    # 将空格替换为特殊字符
    result_text = result_text.replace(' ', '_SPACE_')
    for word in result_text.split(" "):
        # 在输出时将特殊字符替换回空格
        yield word.replace('_SPACE_', ' ')
        time.sleep(0.02)


# 定义⼀个聊天记录到会话管理中
if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "ai", "content": "你好，我是逆蝶，你有什么想说的"}]
    st.session_state["tongyi"] = Tongyi_llm()
    st.session_state["tongyi"].create_llm(key)
    st.session_state["chain"] = st.session_state["tongyi"].get_chain()

print(st.session_state["messages"])

# 文件上传
uploaded_files = st.file_uploader("选择一个或多个文件", accept_multiple_files=True, type=["txt", "pdf", "csv"])

if uploaded_files:
    st.write('你选择了以下文件：')
    for uploaded_file in uploaded_files:
        st.write(uploaded_file.name)

for message in st.session_state["messages"]:
    with st.chat_message(message["role"]):
        st.write(message["content"])

user_input = st.chat_input("我在看着你👁️")

if user_input:
    st.session_state["messages"].append({"role": "human", "content": user_input})
    with st.chat_message("human"):
        st.write_stream(stream_data(user_input))

if user_input:
    if not key:
        st.info("请输⼊你的密钥")
        st.stop()
    with st.spinner("AI正在思考🥵🥵🥵🥵🥵"):
        if uploaded_files:
            response, relvt = get_response_from_file(key, uploaded_files, user_input, chain_choice, strategy_choice)
            result = "请根据以下信息" + response + "回答" + user_input + "你的回答开头为：根据您提供的文件，[你的内容]"
            message = st.session_state["tongyi"].invoke(user_input=result)
        else:
            message = st.session_state["tongyi"].invoke(user_input=user_input)
        st.session_state["messages"].append({"role": "ai", "content": message})
        with st.chat_message("ai"):
            st.write_stream(stream_data(message))
        if uploaded_files:
            with st.expander("相似片段"):
                for relvt_cont in relvt:
                    st.write(relvt_cont)
